// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

// Internal metadata serialization matters

#pragma once

#include <cstdint>
#include <cstring>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include <flatbuffers/flatbuffers.h>

#include <nebula/core/buffer.h>
#include <nebula/io/type_fwd.h>
#include <nebula/ipc/message.h>

#include <nebula/core/sparse_tensor.h>
#include <turbo/utility/status.h>
#include <nebula/types/type_fwd.h>
#include <turbo/base/macros.h>

#include <nebula/fbs/Message_generated.h>
#include <nebula/fbs/Schema_generated.h>
#include <nebula/fbs/SparseTensor_generated.h>  // IWYU pragma: keep

namespace nebula {

    namespace flatbuf = tech::kumo::nebula::flatbuf;

    namespace ipc {

        class DictionaryFieldMapper;

        class DictionaryMemo;

        namespace internal {

            using KeyValueOffset = flatbuffers::Offset<flatbuf::KeyValue>;
            using KVVector = flatbuffers::Vector<KeyValueOffset>;

            // This 0xFFFFFFFF value is the first 4 bytes of a valid IPC message
            constexpr int32_t kIpcContinuationToken = -1;

            static constexpr flatbuf::MetadataVersion kCurrentMetadataVersion =
                    flatbuf::MetadataVersion::V5;

            static constexpr flatbuf::MetadataVersion kLatestMetadataVersion =
                    flatbuf::MetadataVersion::V5;

            static constexpr flatbuf::MetadataVersion kMinMetadataVersion =
                    flatbuf::MetadataVersion::V4;

            // These functions are used in unit tests
            TURBO_EXPORT
            MetadataVersion GetMetadataVersion(flatbuf::MetadataVersion version);

            TURBO_EXPORT
            flatbuf::MetadataVersion MetadataVersionToFlatbuffer(MetadataVersion version);

            // Whether the type has a validity bitmap in the given IPC version
            bool has_validity_bitmap(Type::type type_id, MetadataVersion version);

            static constexpr const char *kArrowMagicBytes = "ARROW1";

            struct FieldMetadata {
                int64_t length;
                int64_t null_count;
                int64_t offset;
            };

            struct BufferMetadata {
                /// The relative offset into the memory page to the starting byte of the buffer
                int64_t offset;

                /// Absolute length in bytes of the buffer
                int64_t length;
            };

            struct FileBlock {
                int64_t offset;
                int32_t metadata_length;
                int64_t body_length;
            };

            // Low-level utilities to help with reading Flatbuffers data.

#define CHECK_FLATBUFFERS_NOT_NULL(fb_value, name)             \
      if ((fb_value) == nullptr) {                                 \
        return turbo::io_error("Unexpected null field ", name,     \
                               " in flatbuffer-encoded metadata"); \
      }

            template<typename T>
            inline uint32_t FlatBuffersVectorSize(const flatbuffers::Vector<T> *vec) {
                return (vec == nullptr) ? 0 : vec->size();
            }

            inline std::string StringFromFlatbuffers(const flatbuffers::String *s) {
                return (s == nullptr) ? "" : s->str();
            }

            // read interface classes. We do not fully deserialize the flatbuffers so that
            // individual fields metadata can be retrieved from very large schema without
            //

            // Construct a complete Schema from the message and add
            // dictionary-encoded fields to a DictionaryMemo instance. May be
            // expensive for very large schemas if you are only interested in a
            // few fields
            TURBO_EXPORT
            turbo::Status get_schema(const void *opaque_schema, DictionaryMemo *dictionary_memo,
                             std::shared_ptr<Schema> *out);

            TURBO_EXPORT
            turbo::Status GetTensorMetadata(BufferSpan metadata, std::shared_ptr<DataType> *type,
                                     std::vector<int64_t> *shape, std::vector<int64_t> *strides,
                                     std::vector<std::string> *dim_names);

            // EXPERIMENTAL: Extracting metadata of a SparseCOOIndex from the message
            TURBO_EXPORT
            turbo::Status GetSparseCOOIndexMetadata(const flatbuf::SparseTensorIndexCOO *sparse_index,
                                             std::shared_ptr<DataType> *indices_type);

            // EXPERIMENTAL: Extracting metadata of a SparseCSXIndex from the message
            TURBO_EXPORT
            turbo::Status GetSparseCSXIndexMetadata(const flatbuf::SparseMatrixIndexCSX *sparse_index,
                                             std::shared_ptr<DataType> *indptr_type,
                                             std::shared_ptr<DataType> *indices_type);

            // EXPERIMENTAL: Extracting metadata of a SparseCSFIndex from the message
            TURBO_EXPORT
            turbo::Status GetSparseCSFIndexMetadata(const flatbuf::SparseTensorIndexCSF *sparse_index,
                                             std::vector<int64_t> *axis_order,
                                             std::vector<int64_t> *indices_size,
                                             std::shared_ptr<DataType> *indptr_type,
                                             std::shared_ptr<DataType> *indices_type);

            // EXPERIMENTAL: Extracting metadata of a sparse tensor from the message
            TURBO_EXPORT
            turbo::Status GetSparseTensorMetadata(BufferSpan metadata, std::shared_ptr<DataType> *type,
                                           std::vector<int64_t> *shape,
                                           std::vector<std::string> *dim_names, int64_t *length,
                                           SparseTensorFormat::type *sparse_tensor_format_id);

            TURBO_EXPORT
            turbo::Status GetKeyValueMetadata(const KVVector *fb_metadata,
                                       std::shared_ptr<KeyValueMetadata> *out);

            TURBO_EXPORT
            turbo::Status ConcreteTypeFromFlatbuffer(flatbuf::Type type, const void *type_data,
                                              FieldVector children, std::shared_ptr<DataType> *out);

            template<typename RootType>
            bool VerifyFlatbuffers(const uint8_t *data, int64_t size) {
                // Heuristic: tables in a Nebula flatbuffers buffer must take at least 1 bit
                // each in average (ARROW-11559).
                // Especially, the only recursive table (the `Field` table in Schema.fbs)
                // must have a non-empty `type` member.
                flatbuffers::Verifier verifier(
                        data, static_cast<size_t>(size),
                        /*max_depth=*/128,
                        /*max_tables=*/static_cast<flatbuffers::uoffset_t>(8 * size));
                return verifier.VerifyBuffer<RootType>(nullptr);
            }

            static inline turbo::Status VerifyMessage(const uint8_t *data, int64_t size,
                                               const flatbuf::Message **out) {
                if (!VerifyFlatbuffers<flatbuf::Message>(data, size)) {
                    return turbo::io_error("Invalid flatbuffers message.");
                }
                *out = flatbuf::GetMessage(data);
                return turbo::OkStatus();
            }

            // serialize nebula::Schema as a Flatbuffer
            TURBO_EXPORT
            turbo::Status WriteSchemaMessage(const Schema &schema, const DictionaryFieldMapper &mapper,
                                      const IpcWriteOptions &options, std::shared_ptr<Buffer> *out);

            // This function is used in a unit test
            TURBO_EXPORT
            turbo::Status WriteRecordBatchMessage(
                    const int64_t length, const int64_t body_length,
                    const std::shared_ptr<const KeyValueMetadata> &custom_metadata,
                    const std::vector<FieldMetadata> &nodes, const std::vector<BufferMetadata> &buffers,
                    const std::vector<int64_t> &variadic_counts, const IpcWriteOptions &options,
                    std::shared_ptr<Buffer> *out);

            TURBO_EXPORT
            turbo::Result<std::shared_ptr<Buffer>> WriteTensorMessage(const Tensor &tensor,
                                                               const int64_t buffer_start_offset,
                                                               const IpcWriteOptions &options);

            TURBO_EXPORT
            turbo::Result<std::shared_ptr<Buffer>> WriteSparseTensorMessage(
                    const SparseTensor &sparse_tensor, int64_t body_length,
                    const std::vector<BufferMetadata> &buffers, const IpcWriteOptions &options);

            TURBO_EXPORT
            turbo::Status WriteFileFooter(const Schema &schema, const std::vector<FileBlock> &dictionaries,
                                   const std::vector<FileBlock> &record_batches,
                                   const std::shared_ptr<const KeyValueMetadata> &metadata,
                                   io::OutputStream *out);

            TURBO_EXPORT
            turbo::Status WriteDictionaryMessage(
                    const int64_t id, const bool is_delta, const int64_t length,
                    const int64_t body_length,
                    const std::shared_ptr<const KeyValueMetadata> &custom_metadata,
                    const std::vector<FieldMetadata> &nodes, const std::vector<BufferMetadata> &buffers,
                    const std::vector<int64_t> &variadic_counts, const IpcWriteOptions &options,
                    std::shared_ptr<Buffer> *out);

            static inline turbo::Result<std::shared_ptr<Buffer>> WriteFlatbufferBuilder(
                    flatbuffers::FlatBufferBuilder &fbb,  // NOLINT non-const reference
                    MemoryPool *pool = default_memory_pool()) {
                int32_t size = fbb.GetSize();

                TURBO_MOVE_OR_RAISE(auto result, allocate_buffer(size, pool));

                uint8_t *dst = result->mutable_data();
                memcpy(dst, fbb.GetBufferPointer(), size);
                // R build with openSUSE155 requires an explicit shared_ptr construction
                return std::shared_ptr<Buffer>(std::move(result));
            }

            TURBO_EXPORT
            flatbuf::TimeUnit ToFlatbufferUnit(TimeUnit::type unit);

            TURBO_EXPORT
            TimeUnit::type FromFlatbufferUnit(flatbuf::TimeUnit unit);

        }  // namespace internal
    }  // namespace ipc
}  // namespace nebula
